﻿using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CodeActions;
using Microsoft.CodeAnalysis.CodeRefactorings;
using Microsoft.CodeAnalysis.CSharp;
using Microsoft.CodeAnalysis.CSharp.Syntax;
using Microsoft.CodeAnalysis.Formatting;
using Microsoft.CodeAnalysis.Simplification;
using Microsoft.CodeAnalysis.Text;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;

namespace CodeAnalysisRuleCollection.Maths.CSharp
{
    [ExportCodeRefactoringProvider("MathPower", LanguageNames.CSharp)]
    internal class MathPow : ICodeRefactoringProvider
    {
        public async Task<IEnumerable<CodeAction>> GetRefactoringsAsync(Document document, TextSpan span, CancellationToken cancellationToken)
        {
            var root = await document.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(false) as CompilationUnitSyntax;
            var model = await document.GetSemanticModelAsync(cancellationToken).ConfigureAwait(false);

            IList<CodeAction> actions = new List<CodeAction>();

            if ((var invocation = root.FindToken(span.Start).Parent.AncestorsAndSelf().OfType<InvocationExpressionSyntax>().FirstOrDefault()) != null)
            {
                // Math.Pow
                if (
                    IsCallToMathPow(model, invocation, cancellationToken) &&
                    invocation.ArgumentList.Arguments.Count == 2
                )
                {
                    Optional<object> @base = model.GetConstantValue(invocation.ArgumentList.Arguments[0].Expression, cancellationToken);
                    if (@base.HasValue)
                    {
                        if (@base.Value.Equals(0))
                            actions.Add(CodeAction.Create("Simplify call to Math.Pow() to 0", c => ChangeCallToValueAsync(document, invocation, 0, model, c)));

                        else if (@base.Value.Equals(1))
                            actions.Add(CodeAction.Create("Simplify call to Math.Pow() to 1", c => ChangeCallToValueAsync(document, invocation, 1, model, c)));
                    }

                    Optional<object> exponent = model.GetConstantValue(invocation.ArgumentList.Arguments[1].Expression, cancellationToken);
                    if (exponent.HasValue)
                    {
                        if (exponent.Value.Equals(0))
                            actions.Add(CodeAction.Create("Simplify call to Math.Pow()", c => ChangeCallToValueAsync(document, invocation, 1, model, c)));

                        else if (exponent.Value.Equals(0.5))
                            actions.Add(CodeAction.Create("Change use of Math.Pow() to Math.Sqrt()", c => ChangePowToSqrtAsync(document, invocation, model, c)));

                        else if (exponent.Value.Equals(1))
                            actions.Add(CodeAction.Create("Simplify call to Math.Pow()", c => SimplifyPowOneAsync(document, invocation, model, c)));

                        else if (exponent.Value.Equals(2))
                            actions.Add(CodeAction.Create("Change use of Math.Pow() to multiplication", c => ChangePowToMultiplicationAsync(document, invocation, model, c)));
                    }

                    // nested calls
                    if ((var nested = invocation.ArgumentList.Arguments[0].Expression as InvocationExpressionSyntax) != null)
                    {
                        if (IsCallToMathSqrt(model, nested, cancellationToken))
                            actions.Add(CodeAction.Create("Simplify nested powers", c => SimplifySqrtInPowAsync(document, invocation, nested, model, c)));

                        else if (IsCallToMathPow(model, nested, cancellationToken))
                            actions.Add(CodeAction.Create("Simplify nested powers", c => SimplifyPowInPowAsync(document, invocation, nested, model, c)));
                    }
                }

                // Math.Sqrt
                else if (
                    IsCallToMathSqrt(model, invocation, cancellationToken) &&
                    invocation.ArgumentList.Arguments.Count == 1
                )
                {
                    actions.Add(CodeAction.Create("Change use of Math.Sqrt() to Math.Pow()", c => ChangeSqrtToPowAsync(document, invocation, model, c)));

                    Optional<object> @base = model.GetConstantValue(invocation.ArgumentList.Arguments[0].Expression, cancellationToken);
                    if (@base.HasValue)
                    {
                        if (@base.Value.Equals(0))
                            actions.Add(CodeAction.Create("Simplify call to Math.Sqrt() to 0", c => ChangeCallToValueAsync(document, invocation, 0, model, c)));

                        else if (@base.Value.Equals(1))
                            actions.Add(CodeAction.Create("Simplify call to Math.Sqrt() to 1", c => ChangeCallToValueAsync(document, invocation, 1, model, c)));
                    }

                    // nested calls
                    if ((var nested = invocation.ArgumentList.Arguments[0].Expression as InvocationExpressionSyntax) != null)
                    {
                        if (IsCallToMathSqrt(model, nested, cancellationToken))
                            actions.Add(CodeAction.Create("Simplify nested powers", c => SimplifySqrtInSqrtAsync(document, invocation, nested, model, c)));

                        else if (IsCallToMathPow(model, nested, cancellationToken))
                            actions.Add(CodeAction.Create("Simplify nested powers", c => SimplifyPowInSqrtAsync(document, invocation, nested, model, c)));
                    }
                }
            }

            // Multiplication
            if ((var binary = root.FindToken(span.Start).Parent.AncestorsAndSelf().OfType<BinaryExpressionSyntax>().FirstOrDefault()) != null)
            {
                if (binary.IsKind(SyntaxKind.MultiplyExpression) && SyntaxFactory.AreEquivalent(binary.Left, binary.Right))
                {
                    actions.Add(CodeAction.Create("Change multiplication to use of Math.Pow()", c => ChangeMultiplicationToPowAsync(document, binary, model, c)));

                    if (
                        (var invocation = binary.Left as InvocationExpressionSyntax) != null &&
                        IsCallToMathSqrt(model, invocation, cancellationToken)
                    )
                        actions.Add(CodeAction.Create("Simplify square of square roots", c => SimpliySqrtTimesSqrtAsync(document, binary, invocation, model, c)));
                }
            }

            return actions;
        }

        private async Task<Document> ChangePowToMultiplicationAsync(Document document, InvocationExpressionSyntax invocation, SemanticModel model, CancellationToken c)
        {
            var root = await document.GetSyntaxRootAsync(c);

            // replace node
            var newRoot = root.ReplaceNode<SyntaxNode, ExpressionSyntax>(invocation,
                SyntaxFactory.BinaryExpression(SyntaxKind.MultiplyExpression,
                    invocation.ArgumentList.Arguments[0].Expression,
                    invocation.ArgumentList.Arguments[0].Expression
                )
                    .WithAdditionalAnnotations(Simplifier.Annotation)
                    .WithAdditionalAnnotations(Formatter.Annotation)
            );

            return document.WithSyntaxRoot(newRoot);
        }

        private async Task<Document> ChangePowToSqrtAsync(Document document, InvocationExpressionSyntax invocation, SemanticModel model, CancellationToken c)
        {
            var root = await document.GetSyntaxRootAsync(c);

            // replace node
            var newRoot = root.ReplaceNode(invocation,
                SyntaxFactory.InvocationExpression(
                    SyntaxFactory.MemberAccessExpression(SyntaxKind.SimpleMemberAccessExpression, SyntaxFactory.IdentifierName("Math"), SyntaxFactory.IdentifierName("Sqrt")),
                    SyntaxFactory.ArgumentList(
                        SyntaxFactory.SeparatedList(new[] {
                            invocation.ArgumentList.Arguments[0]
                        })
                    )
                )
                    .WithAdditionalAnnotations(Simplifier.Annotation)
            );

            return document.WithSyntaxRoot(newRoot);
        }

        private async Task<Document> ChangeSqrtToPowAsync(Document document, InvocationExpressionSyntax invocation, SemanticModel model, CancellationToken c)
        {
            var root = await document.GetSyntaxRootAsync(c);

            // replace node
            var newRoot = root.ReplaceNode(invocation,
                SyntaxFactory.InvocationExpression(
                    SyntaxFactory.MemberAccessExpression(SyntaxKind.SimpleMemberAccessExpression, SyntaxFactory.IdentifierName("Math"), SyntaxFactory.IdentifierName("Pow")),
                    SyntaxFactory.ArgumentList(
                        SyntaxFactory.SeparatedList(new[] {
                            invocation.ArgumentList.Arguments[0],
                            SyntaxFactory.Argument(SyntaxFactory.LiteralExpression(SyntaxKind.NumericLiteralExpression, SyntaxFactory.Literal(0.5)))
                        })
                    )
                )
                    .WithAdditionalAnnotations(Formatter.Annotation)
            );

            return document.WithSyntaxRoot(newRoot);
        }

        private async Task<Document> ChangeMultiplicationToPowAsync(Document document, BinaryExpressionSyntax binary, SemanticModel model, CancellationToken c)
        {
            var root = await document.GetSyntaxRootAsync(c);

            // replace node
            var newRoot = root.ReplaceNode<SyntaxNode, ExpressionSyntax>(binary,
                SyntaxFactory.InvocationExpression(
                    SyntaxFactory.MemberAccessExpression(SyntaxKind.SimpleMemberAccessExpression, SyntaxFactory.IdentifierName("Math"), SyntaxFactory.IdentifierName("Pow")),
                    SyntaxFactory.ArgumentList(
                        SyntaxFactory.SeparatedList(new[] {
                            SyntaxFactory.Argument(binary.Left),
                            SyntaxFactory.Argument(SyntaxFactory.LiteralExpression(SyntaxKind.NumericLiteralExpression, SyntaxFactory.Literal(2)))
                        })
                    )
                )
                    .WithAdditionalAnnotations(Formatter.Annotation)
            );

            return document.WithSyntaxRoot(newRoot);
        }

        private async Task<Document> ChangeCallToValueAsync(Document document, InvocationExpressionSyntax invocation, int value, SemanticModel model, CancellationToken c)
        {
            var root = await document.GetSyntaxRootAsync(c);

            // replace node
            var newRoot = root.ReplaceNode<SyntaxNode, ExpressionSyntax>(invocation,
                SyntaxFactory.LiteralExpression(SyntaxKind.NumericLiteralExpression, SyntaxFactory.Literal(value))
                    .WithAdditionalAnnotations(Formatter.Annotation)
            );

            return document.WithSyntaxRoot(newRoot);
        }

        private async Task<Document> SimpliySqrtTimesSqrtAsync(Document document, BinaryExpressionSyntax binary, InvocationExpressionSyntax invocation, SemanticModel model, CancellationToken c)
        {
            var root = await document.GetSyntaxRootAsync(c);

            // replace node
            var newRoot = root.ReplaceNode(binary,
                invocation.ArgumentList.Arguments[0].Expression
                    .WithAdditionalAnnotations(Formatter.Annotation)
            );

            return document.WithSyntaxRoot(newRoot);
        }

        private async Task<Document> SimplifyPowOneAsync(Document document, InvocationExpressionSyntax invocation, SemanticModel model, CancellationToken c)
        {
            var root = await document.GetSyntaxRootAsync(c);

            // replace node
            var newRoot = root.ReplaceNode<SyntaxNode, ExpressionSyntax>(invocation,
                invocation.ArgumentList.Arguments[0].Expression
            );

            return document.WithSyntaxRoot(newRoot);
        }

        private async Task<Document> SimplifySqrtInPowAsync(Document document, InvocationExpressionSyntax outer, InvocationExpressionSyntax inner, SemanticModel model, CancellationToken c)
        {
            var root = await document.GetSyntaxRootAsync(c);

            // replace node
            var newRoot = root.ReplaceNode(outer,
                outer
                    .WithArgumentList(
                        SyntaxFactory.ArgumentList(
                            SyntaxFactory.SeparatedList(new[] {
                                inner.ArgumentList.Arguments[0],
                                SyntaxFactory.Argument(
                                    SyntaxFactory.BinaryExpression(SyntaxKind.MultiplyExpression,
                                        outer.ArgumentList.Arguments[1].Expression,
                                        SyntaxFactory.LiteralExpression(SyntaxKind.NumericLiteralExpression, SyntaxFactory.Literal(0.5))
                                    )
                                )
                            })
                        )
                    )
                    .WithAdditionalAnnotations(Formatter.Annotation)
            );

            return document.WithSyntaxRoot(newRoot);
        }

        private async Task<Document> SimplifyPowInPowAsync(Document document, InvocationExpressionSyntax outer, InvocationExpressionSyntax inner, SemanticModel model, CancellationToken c)
        {
            var root = await document.GetSyntaxRootAsync(c);

            // replace node
            var newRoot = root.ReplaceNode(outer,
                outer
                    .WithArgumentList(
                        SyntaxFactory.ArgumentList(
                            SyntaxFactory.SeparatedList(new[] {
                                inner.ArgumentList.Arguments[0],
                                SyntaxFactory.Argument(
                                    SyntaxFactory.BinaryExpression(SyntaxKind.MultiplyExpression,
                                        inner.ArgumentList.Arguments[1].Expression,
                                        outer.ArgumentList.Arguments[1].Expression
                                    )
                                )
                            })
                        )
                    )
                    .WithAdditionalAnnotations(Formatter.Annotation)
            );

            return document.WithSyntaxRoot(newRoot);
        }

        private async Task<Document> SimplifyPowInSqrtAsync(Document document, InvocationExpressionSyntax outer, InvocationExpressionSyntax inner, SemanticModel model, CancellationToken c)
        {
            var root = await document.GetSyntaxRootAsync(c);

            // replace node
            var newRoot = root.ReplaceNode(outer,
                inner
                    .WithArgumentList(
                        SyntaxFactory.ArgumentList(
                            SyntaxFactory.SeparatedList(new[] {
                                inner.ArgumentList.Arguments[0],
                                SyntaxFactory.Argument(
                                    SyntaxFactory.BinaryExpression(SyntaxKind.MultiplyExpression,
                                        inner.ArgumentList.Arguments[1].Expression,
                                        SyntaxFactory.LiteralExpression(SyntaxKind.NumericLiteralExpression, SyntaxFactory.Literal(0.5))
                                    )
                                )
                            })
                        )
                    )
                    .WithAdditionalAnnotations(Formatter.Annotation)
            );

            return document.WithSyntaxRoot(newRoot);
        }

        private async Task<Document> SimplifySqrtInSqrtAsync(Document document, InvocationExpressionSyntax outer, InvocationExpressionSyntax inner, SemanticModel model, CancellationToken c)
        {
            var root = await document.GetSyntaxRootAsync(c);

            // replace node
            var newRoot = root.ReplaceNode(outer,
                SyntaxFactory.InvocationExpression(
                    SyntaxFactory.MemberAccessExpression(SyntaxKind.SimpleMemberAccessExpression, SyntaxFactory.IdentifierName("Math"), SyntaxFactory.IdentifierName("Pow")),
                    SyntaxFactory.ArgumentList(
                        SyntaxFactory.SeparatedList(new[] {
                            inner.ArgumentList.Arguments[0],
                            SyntaxFactory.Argument(SyntaxFactory.LiteralExpression(SyntaxKind.NumericLiteralExpression, SyntaxFactory.Literal(0.25)))
                        })
                    )
                )
                    .WithAdditionalAnnotations(Formatter.Annotation)
            );

            return document.WithSyntaxRoot(newRoot);
        }

        private bool IsCallToMathPow(SemanticModel semanticModel, InvocationExpressionSyntax invocation, CancellationToken cancellationToken)
        {
            var call = invocation.Expression as MemberAccessExpressionSyntax;
            if (call == null)
                return false;

            var methodSymbol = semanticModel.GetSymbolInfo(call, cancellationToken).Symbol
                as IMethodSymbol;

            if (methodSymbol == null ||
                methodSymbol.Name != "Pow")
                return false;

            var math = semanticModel.Compilation.GetTypeByMetadataName(typeof(Math).FullName);

            if (math == null ||
                !methodSymbol.ContainingType.Equals(math))
                return false;

            return true;
        }

        private bool IsCallToMathSqrt(SemanticModel semanticModel, InvocationExpressionSyntax invocation, CancellationToken cancellationToken)
        {
            var call = invocation.Expression as MemberAccessExpressionSyntax;
            if (call == null)
                return false;

            var methodSymbol = semanticModel.GetSymbolInfo(call, cancellationToken).Symbol
                as IMethodSymbol;

            if (methodSymbol == null ||
                methodSymbol.Name != "Sqrt")
                return false;

            var math = semanticModel.Compilation.GetTypeByMetadataName(typeof(Math).FullName);

            if (math == null ||
                !methodSymbol.ContainingType.Equals(math))
                return false;

            return true;
        }
    }
}
